package com.shujia

import com.shujia.useraction.UserActionMerge
import com.shujia.utils.{Constants, SparkMain, SparkTool}
import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.{ALS, ALSModel}
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.types.{FloatType, IntegerType}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import redis.clients.jedis.Jedis

object UserRecOnALS extends SparkMain {
  override def run(day: String): Unit = {

    val spark: SparkSession = SparkTool.getSparkSession("UserRecOnALS")
    import spark.implicits._
    import org.apache.spark.sql.functions._


    val mergeDF: DataFrame = UserActionMerge(day).load()


    val actionDF: DataFrame = mergeDF
      .where($"uid" =!= "null")
      .withColumn("userId", rank() over Window.orderBy("uid"))
      .select($"uid", $"userId", $"itemId".cast(IntegerType), $"num".cast(FloatType))
    // 将数据分为训练集和测试集
    val Array(training, test): Array[Dataset[Row]] = actionDF.randomSplit(Array(0.8, 0.2))

    // 构建ALS模型
    val als: ALS = new ALS()
      .setMaxIter(10) // 设置最大的迭代次数
      .setRegParam(0.01) // 设置缩放比例
      .setUserCol("userId") // 用户id列
      .setItemCol("itemId") // 物品id列
      .setRatingCol("num") // 评分列

    // 使用训练集训练模型
    val model: ALSModel = als.fit(training)

    // 将预测值为空的删除
    model.setColdStartStrategy("drop")

    // 使用测试集测试模型
    val predictions: DataFrame = model.transform(test)

    //    predictions.show()
    // 使用rmse均方根误差来评估模型
    val evaluator: RegressionEvaluator = new RegressionEvaluator()
      .setMetricName("rmse") // 指定指标的名称
      .setLabelCol("num") // 指定label标签列
      .setPredictionCol("prediction") // 指定ALS预测结果列
    val rmse: Double = evaluator.evaluate(predictions)
    println(s"Root-mean-square error = $rmse")


    // 对每个用户推荐10部影片
    val users: Dataset[Row] = actionDF.select(als.getUserCol).distinct()
    val userSubsetRecs: DataFrame = model.recommendForUserSubset(users, 10)

    //    userSubsetRecs.show(1000, truncate = false)

    /**
     * root
     * |-- userId: integer (nullable = false)
     * |-- recommendations: array (nullable = true)
     * |    |-- element: struct (containsNull = true)
     * |    |    |-- itemId: integer (nullable = true)
     * |    |    |-- rating: float (nullable = true)
     */
    //    userSubsetRecs.printSchema()

    val recDF: DataFrame = userSubsetRecs
      .join(actionDF, "userId")
      .select($"uid", explode($"recommendations") as "recList")
      .select($"uid", $"recList.itemId" as "itemId", $"recList.rating" as "rating")


    recDF.foreachPartition(iter => {
      // 创建Redis连接
      val jedis: Jedis = new Jedis(Constants.REDIS_HOST, Constants.REDIS_PORT)
      iter.foreach {
        case Row(uid, itemId, rating) =>
          jedis.hset(s"rec_list:$uid", itemId.toString, rating.toString)
          jedis.expire(s"rec_list:$uid", 24 * 60 * 60)

      }
    })

  }

}
